import math
import os
import pickle

import numpy as np
import trimesh
import trimesh.creation
# from model.unit.Function.constraint import *
from environment.unit.base.ClampObject import ClampBlock
from trimesh.collision import CollisionManager
from trimesh.points import point_plane_distance
from trimesh.transformations import concatenate_matrices
from trimesh.transformations import identity_matrix
from trimesh.transformations import rotation_matrix, translation_matrix
from trimesh.transformations import transform_points
from trimesh.transformations import unit_vector


class Shape(object):
    AXIS_X = (np.array([0, 0, 0], dtype=float), np.array([1, 0, 0], dtype=float))
    AXIS_Y = (np.array([0, 0, 0], dtype=float), np.array([0, 1, 0], dtype=float))
    AXIS_Z = (np.array([0, 0, 0], dtype=float), np.array([0, 0, 1], dtype=float))

    def __init__(self, vs=None, transform=None, face_colors=None):
        self.object_list = []
        self.name_list = []
        self.constrain_list = []
        self.relevance_list = []
        self.is_hidden = False
        self.name = 'CL'
        self.type = 'None'
        self.category = ''
        self.description = ''
        self.top_group = self
        self.is_fixed = False
        self.transform_mode = 0
        self.feature_point_list = []
        self.assembly_area_list = []
        self.assembly_center_point_list = []
        self.stl_path = ''
        self.clamp_block_list = []
        self.top_area = None
        self.bottom_area = None
        if vs is None:
            vs = {}

        if transform is None:
            transform = identity_matrix()
        self.transform = transform
        self.vs = vs
        self.mesh = None
        self.face_colors = face_colors
        self._transform_save = None

    def __repr__(self):
        return 'Shape'

    def __add__(self, other):
        return Group.create_group(self, other)

    @staticmethod
    def get_direction(position):
        assert position[0] in '+-'
        reverse = position[0] == '-'
        assert position[1] in "xyz"

        if position[1] == 'x':
            direction = np.array([1, 0, 0], dtype=float)
        elif position[1] == 'y':
            direction = np.array([0, 1, 0], dtype=float)
        elif position[1] == 'z':
            direction = np.array([0, 0, 1], dtype=float)
        else:
            raise NotImplementedError()

        if reverse:
            direction *= -1

        return direction

    def min_distance(self, env_cm, transform=None, **kwargs):
        if isinstance(env_cm, (trimesh.Trimesh, Shape)):
            if isinstance(env_cm, Shape):
                env_mesh = env_cm.get_mesh()
            else:
                env_mesh = env_cm
            env_cm = CollisionManager()
            env_cm.add_object("env", env_mesh)

        ret = env_cm.min_distance_single(self.get_mesh(), transform=transform, **kwargs)
        return ret

    def is_collision(self, env_mesh):
        c = trimesh.collision.CollisionManager()
        c.add_object('env', env_mesh)
        c.add_object('box', self.get_mesh())

        return c.in_collision_internal(return_names=False, return_data=False)

    def get_z_space(self, env_mesh, direction, gap_step=2, offset_max=400):
        c = trimesh.collision.CollisionManager()
        c.add_object('env', env_mesh)

        step = 0
        while True:
            c.add_object('box', self.get_mesh())
            if not c.in_collision_internal(return_names=False, return_data=False):
                ans = step * gap_step
                break
            else:
                c.remove_object('box')
            if direction == 'up':
                self.move([0, 0, gap_step])
            elif direction == 'down':
                self.move([0, 0, -gap_step])
            step += 1
            if step > offset_max:
                assert 0

        return ans

    def collision_region(self, env_cm, direction=None, limit=1000):
        if direction is None:
            direction = [0, 0, 1]
        direction = unit_vector(np.asarray(direction, float))

        mesh = self.get_mesh()

        arr = []
        status = False
        for i in range(-limit, limit):
            distance = env_cm.min_distance_single(mesh, transform=translation_matrix(direction * i))
            if status is False:
                if distance < 0.5:
                    arr.append(i)
                    status = True
            else:
                if distance > 0.5:
                    arr.append(i)
                    status = False

        return arr

    # def end_constrain(self):
    #     delete_all_constrain(self)

    @staticmethod
    def export_assembly_point_stl(box, path):
        radius = min(box.vs['ly'], box.vs['lz']) / 20
        assembly_point_list = box.get_all_assembly_point()
        for assembly_center in assembly_point_list:
            assembly_point = assembly_center.point
            assembly_point.construct()
            assembly_mark = trimesh.creation.uv_sphere(radius)
            assembly_mark.apply_transform(translation_matrix(assembly_point.origin_point))
            assembly_mark.apply_transform(assembly_point.reference_box.transform)
            assembly_mark.export(path + r'/' + assembly_center.name + '.stl', file_type='stl')

    @staticmethod
    def export_xyz_point_stl(box, path):
        radius = min(box.vs['ly'], box.vs['lz']) / 20

        x_mark = trimesh.creation.uv_sphere(radius)
        x_mark.apply_transform(translation_matrix([box.vs['lx'] * 0.5, 0, 0]))
        x_mark.apply_transform(box.transform)
        x_mark.export(path + r'/axis_x.stl', file_type='stl')

        y_mark = trimesh.creation.uv_sphere(radius)
        y_mark.apply_transform(translation_matrix([0, box.vs['ly'] * 0.5, 0]))
        y_mark.apply_transform(box.transform)
        y_mark.export(path + r'/axis_y.stl', file_type='stl')

        z_mark = trimesh.creation.uv_sphere(radius)
        z_mark.apply_transform(translation_matrix([0, 0, box.vs['lz'] * 0.5]))
        z_mark.apply_transform(box.transform)
        z_mark.export(path + r'/axis_z.stl', file_type='stl')
        return [path + r'/axis_x.stl', path + r'/axis_y.stl', path + r'/axis_z.stl']

    @staticmethod
    def save_obj(obj, file_path='data.pkl'):
        with open(file_path, 'wb') as f:
            pickle.dump(obj, f)

    def save_stl(self, file_path='temp', mesh=None):
        obj = self.get_mesh()
        if file_path is None:
            file_path = 'temp' + '/' + self.name + '.stl'
        if not file_path.endswith('.stl'):
            file_path = file_path + '/' + self.name + '.stl'
        if os.path.exists(file_path):
            os.remove(file_path)
        obj.export(file_path, file_type='stl')
        return os.path.abspath(file_path)

    @staticmethod
    def load_obj(file_path='data.pkl'):
        with open(file_path, 'rb') as f:
            obj = pickle.load(f)
        return obj

    @staticmethod
    def point_vector(start_point, end_point, reserve='xyz'):
        p1 = np.array(start_point, dtype=float)
        p2 = np.array(end_point, dtype=float)
        if 'x' not in reserve:
            p1[0] = 0
            p2[0] = 0
        if 'y' not in reserve:
            p1[1] = 0
            p2[1] = 0
        if 'z' not in reserve:
            p1[2] = 0
            p2[2] = 0

        return p2 - p1

    @staticmethod
    def point_distance(start_point, end_point, reserve='xyz'):
        vec = Shape.point_vector(start_point, end_point, reserve)

        return np.linalg.norm(vec)

    @staticmethod
    def point_2d_distance(point1, point2):
        return math.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2)

    @staticmethod
    def point_z_distance(point1, point2):
        return abs(point1[2] - point2[2])

    @staticmethod
    def point2plane_distance(point, plane):
        return point_plane_distance(point, plane[1], plane[0])

    @staticmethod
    def point_angle(start_point, end_point, reserve='xy'):
        assert reserve == 'xy'
        vec = Shape.point_vector(start_point, end_point, reserve)
        return np.rad2deg(np.arctan2(vec[1], vec[0]))

    def construct(self, vs=None):
        if vs is None:
            vs = {}

        self.vs.update(vs)
        self._transform_save = self.transform
        self.transform = identity_matrix()

    def end_construct(self):
        self.transform = identity_matrix()
        self.transform_mode = 2
        self.add_transform(self._transform_save)
        # for obj in self.object_list:
        #     obj.add_transform(self.transform)

    # def is_collision(self, obj):
    #     m = trimesh.collision.CollisionManager()
    #     m.add_object('box1', self.get_mesh())
    #     m.add_object('box2', obj.get_mesh())
    #     obj_mesh = obj.get_mesh()
    #     # print(np.array(self.get_point()))
    #     # print(obj_mesh.contains(np.array(self.get_point())))
    #     distance = m.min_distance_internal(return_names=False, return_data=True)
    #     # print(distance)
    #     is_collision, names, data = m.in_collision_internal(return_names=True, return_data=True)
    #     if m.in_collision_internal(return_names=False, return_data=True):
    #         return True
    #     else:
    #         return False

    def transformed_point(self, point):
        point = transform_points(
            [point],
            self.transform
        )
        return point[0]

    def get_mesh(self):
        raise NotImplementedError()

    def is_direct_relevance(self, obj):
        direct_relevance_list = []
        for constrain in self.constrain_list:
            if constrain[-1] == 2:
                if isinstance(constrain[0], list):
                    part1 = constrain[0][0].reference_object
                    part2 = constrain[1][0].reference_object
                else:
                    part1 = constrain[0].reference_object
                    part2 = constrain[1].reference_object
                if self.name in [part1.name, part2.name]:
                    if self.name == part1.name:
                        obj_part = part2
                    else:
                        obj_part = part1
                    if obj_part not in direct_relevance_list:
                        direct_relevance_list.append(obj_part)
        if obj in direct_relevance_list:
            return True
        else:
            return False

    def is_indirect_relevance(self, obj):
        if self.is_direct_relevance(obj):
            return False
        elif self == obj:
            return False
        else:
            relevance_object_list = self.get_relevance_object()
            if obj in relevance_object_list:
                return True
            else:
                return False

    def is_relevance(self, obj):
        if self == obj:
            return False
        elif self.is_direct_relevance(obj):
            return True
        elif self.is_indirect_relevance(obj):
            return True
        else:
            return False

    def get_direct_relevance_object(self):
        direct_relevance_list = []
        for obj in self.top_group.object_list:
            if self.is_direct_relevance(obj):
                direct_relevance_list.append(obj)
        return direct_relevance_list

    def get_indirect_relevance_object(self):
        indirect_relevance_list = []
        for obj in self.top_group.object_list:
            if self.is_indirect_relevance(obj):
                indirect_relevance_list.append(obj)
        return indirect_relevance_list

    def get_relevance_object(self):
        self.relevance_list = [self]
        self.get_relevance_object_func(self)
        self.relevance_list.remove(self)
        return self.relevance_list

    def get_relevance_object_func(self, origin_obj):
        direct_relevance_list = self.get_direct_relevance_object()
        for obj in direct_relevance_list:
            if obj not in origin_obj.relevance_list:
                origin_obj.relevance_list.append(obj)
                obj.get_relevance_object_func(origin_obj)

    def add_transform(self, transform):
        if self.top_group.transform_mode == 0:  # 约束模式
            # assert not self.is_fixed, 'the Object has been fixed'
            self.transform = concatenate_matrices(transform, self.transform)
            # self.update_relevance()

        elif self.top_group.transform_mode == 1:  # 共动模式
            assert not self.is_fixed, 'the Object has been fixed'
            self.transform = concatenate_matrices(transform, self.transform)
            relevance_obj_list = self.get_relevance_object()
            for relevance_object in relevance_obj_list:
                relevance_object.transform = concatenate_matrices(transform, relevance_object.transform)
        elif self.top_group.transform_mode == 2:  # 自由模式
            self.transform = concatenate_matrices(transform, self.transform)

    def rotate(self, angle, axis):
        """绕轴axis旋转角度angle"""
        origin, direction = axis
        matrix = rotation_matrix(
            np.deg2rad(angle),
            direction,
            origin
        )
        self.add_transform(matrix)

    def move(self, xyz):
        self.add_transform(translation_matrix(xyz))

    def move_point2plane(self, point1, plane1):
        """
        从point1移动到plane1
        """
        point1 = np.asarray(point1, dtype=float)
        origin1, normal1 = plane1
        distance = point_plane_distance(point1, normal1, origin1)
        self.move(-distance * normal1)

    def move_point2point(self, point1, point2):
        """
        从point1移动到point2
        """
        point1 = np.asarray(point1, dtype=float)
        point2 = np.asarray(point2, dtype=float)

        self.move(point2 - point1)

    def get_all_feature_point(self):
        for obj in self.object_list:
            self.feature_point_list += obj.feature_point_list
        return self.feature_point_list

    def get_all_assembly_area(self):
        for obj in self.object_list:
            self.assembly_area_list += obj.assembly_area_list
        return self.assembly_area_list

    def get_all_assembly_point(self):
        for obj in self.object_list:
            self.assembly_center_point_list += obj.assembly_center_point_list
        return self.assembly_center_point_list

    def get_all_clamp_block(self):
        for clamp_block in self.clamp_block_list:
            clamp_block.reference_name = self.name + '/' + clamp_block.obj.name
            clamp_block.construct()
        return self.clamp_block_list

    def get_assembly_area(self, name):
        for object_item in self.object_list:
            for ass_area in object_item.assembly_area_list:
                if ass_area.name == name:
                    return ass_area
        raise IOError('This Assembly_Area Not Exit')

    def get_assembly_center(self, name):
        for object_item in self.object_list:
            for ass_center in object_item.assembly_center_point_list:  # Point()
                if ass_center.name == name:  # 通过名字找到那个装配点
                    ass_center.point.reference_object = self
                    return ass_center
        return None


class Group(Shape):
    def __init__(self, vs=None, transform=None, face_colors=None):
        super().__init__(vs, transform, face_colors)
        self.origin_clamp_block = False
        self.origin_part = None
        self.version = 1.0
        self.construct()

    @staticmethod
    def create_group(*args):
        g = Group()
        for obj in args:
            # obj_temp = obj
            g.add_object(obj, obj.name)
        return g

    def get_moved_point(self, position, xyz):
        position = position.split('.', 1)
        return self.get_object(position[0]).get_moved_point(position[1], xyz)

    def get_point(self, position, offset=None):
        if offset is None:
            offset = [0, 0, 0]
        position = position.split('.', 1)
        return self.get_object(position[0]).get_point(position[1], offset)

    def get_axis(self, position, point=None):
        if '.' not in position:
            direction = self.get_direction(position)
            return self.transformed_point(direction) - self.transformed_point([0, 0, 0])

        position = position.split('.', 1)
        return self.get_object(position[0]).get_axis(position[1], point)

    def get_plane(self, position):
        position = position.split('.', 1)
        return self.get_object(position[0]).get_plane(position[1])

    def get_xy_distance(self, p1, p2):
        return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)


    def add_transform(self, transform):
        super().add_transform(transform)
        for obj in self.object_list:
            obj.add_transform(transform)

    def get_object(self, name):
        if '/' not in name:
            return self.object_list[self.name_list.index(name)]
        else:
            parts = name.split('/', 1)
            obj = self.object_list[self.name_list.index(parts[0])].get_object(parts[1])
            return obj

    def get_object_point(self, position, offset=None):
        if offset is None:
            offset = [0, 0, 0]
        position = position.split('.', 1)
        point1 = self.get_object(position[0]).get_object_point(position[1], offset)
        point1.reference_object = self
        return point1

    def get_object_plane(self, position):
        position = position.split('.', 1)
        plane1 = self.get_object(position[0]).get_object_plane(position[1])
        plane1.reference_object = self
        return plane1

    def get_object_hole_group(self, position):
        position = position.split('.', 1)
        return self.get_object(position[0]).get_object_plane(position[1])

    def get_all_boxes(self, prefix=""):
        boxes = []
        for obj, name in zip(self.object_list, self.name_list):
            if isinstance(obj, Box):
                boxes.append((prefix + name, obj))
            elif isinstance(obj, Group):
                boxes += obj.get_all_boxes(prefix=name + '/')
        return boxes

    def get_all_plane(self, prefix=""):
        planes = []
        boxes = self.get_all_boxes()
        for box in boxes:
            box_name = box[0]
            box_obj = box[1]
            planes += box_obj.get_all_plane(box_name, self)
        return planes

    def __repr__(self):
        dic = {}
        for i in range(len(self.name_list)):
            dic[self.name_list[i]] = str(self.object_list[i])
        return str(dic)

    def del_object(self, name):
        index = self.name_list.index(name)
        del self.object_list[index]
        self.name_list.pop(index)

    def add_object(self, obj, name=None, rps=None):
        # for box in self.get_all_boxes():
        #     box.add_transform(self.transform)

        if name is None:
            if obj.name is None:
                if isinstance(obj, Box) or obj == Box:
                    pre = 'box'
                elif isinstance(obj, Group) or obj == Group:
                    pre = 'group'
                else:
                    pre = 'else'

                name = pre + str(len(self.object_list))
            else:
                name = obj.name

        if isinstance(obj, type):
            vs = {}
            for k in self.vs:
                if k.startswith(name + '/'):
                    vs[k.replace(name + '/', '', 1)] = self.vs[k]

            # print('vs', vs)
            obj = obj(vs)

        if name not in self.name_list:
            self.object_list.append(obj)
            self.name_list.append(name)
        else:
            # vs = self.object_list[self.name_list.index(name)].vs
            # vs.update(obj.vs)
            # obj.vs = vs
            # obj.transform = self.object_list[self.name_list.index(name)].transform
            self.object_list[self.name_list.index(name)] = obj
        obj.name = name
        obj.top_group = self
        if obj.type == 'Block':
            self.clamp_block_list.append(ClampBlock(obj, self, rps=rps))
        else:
            temp_list = obj.get_all_clamp_block()
            for clamp_b in temp_list:
                clamp_b.rps = rps
            self.clamp_block_list.extend(temp_list)
        return obj

    def add_object_list(self, obj_list, obj_name_list=None):
        for index, obj in enumerate(obj_list):
            if obj_name_list is None:
                self.add_object(obj, obj.name)
            else:
                assert len(obj_list) == len(obj_name_list)
                self.add_object(obj, obj_name_list[index])

    def get_mesh(self):
        if len(self.object_list) < 1:
            return None
        self.mesh = None
        for i in range(0, len(self.object_list)):
            if self.object_list[i].is_hidden:
                continue

            sub_mesh = self.object_list[i].get_mesh()
            if sub_mesh is None:
                continue
            if self.mesh is None:
                self.mesh = sub_mesh
            else:
                self.mesh += sub_mesh

        if self.mesh is None:
            return None

        if self.face_colors is not None:
            self.mesh.visual.face_colors = self.face_colors

        return self.mesh


class Cylinder(Shape):
    """ 圆柱 """

    def __init__(self, vs=None, transform=None, face_colors=None):
        if vs is None:
            vs = {}
        if 'lz' not in vs:
            vs['lz'] = 1
        if 'lr' not in vs:
            vs['lr'] = 1

        super().__init__(vs, transform, face_colors)

    def __repr__(self):
        return 'Cylinder'

    def get_point(self, position='', offset=None, world_axis=True):
        if offset is None:
            offset = [0, 0, 0]
        point = np.array([0, 0, 0], dtype=float)
        if '+z' in position:
            point[2] = self.vs['lz'] * 0.5
        elif '-z' in position:
            point[2] = -self.vs['lz'] * 0.5
        point += offset
        if world_axis:
            return self.transformed_point(point)
        else:
            return point

    def get_mesh(self):
        self.mesh = trimesh.creation.cylinder(
            radius=self.vs['lr'],
            height=self.vs['lz']
        )
        self.mesh.apply_transform(self.transform)

        if self.face_colors is not None:
            self.mesh.visual.face_colors = self.face_colors

        return self.mesh


# Cylinder().show()
# sys.exit()


class Box(Shape):
    def __init__(self, vs=None, transform=None, face_colors=None):
        if vs is None:
            vs = {}

        if isinstance(vs, (list, tuple)):
            assert len(vs) == 3
            vs = {
                'lx': vs[0],
                'ly': vs[1],
                'lz': vs[2],
            }

        if 'lx' not in vs:
            vs['lx'] = 1
        if 'ly' not in vs:
            vs['ly'] = 1
        if 'lz' not in vs:
            vs['lz'] = 1
        super().__init__(vs, transform, face_colors)

        # for k, ax in [('rx', self.AXIS_X), ('ry', self.AXIS_Y), ('rz', self.AXIS_Z)]:
        #     if k in self.vs:
        #         self.rotate(self.vs[k], ax)

    def __repr__(self):
        return 'Box'

    def get_moved_point(self, position, xyz, world_axis=True):
        def fun1(plist):
            return np.array([round(v, 3) for v in plist], dtype=float)

        point = np.array([0, 0, 0], dtype=float)
        xyz = np.array(xyz, dtype=float)
        assert '.' not in position

        if '+x' in position:
            point[0] = self.vs['lx'] * 0.5
        elif '-x' in position:
            point[0] = -self.vs['lx'] * 0.5
        if '+y' in position:
            point[1] = self.vs['ly'] * 0.5
        elif '-y' in position:
            point[1] = -self.vs['ly'] * 0.5
        if '+z' in position:
            point[2] = self.vs['lz'] * 0.5
        elif '-z' in position:
            point[2] = -self.vs['lz'] * 0.5
        if position == '':
            point = [0, 0, 0]

        point += xyz
        if world_axis:
            point_world = self.transformed_point(point)
            return point_world
        else:
            return point

    def get_point(self, position='', offset=None, world_axis=True):
        if offset is None:
            offset = [0, 0, 0]
        return self.get_moved_point(position, offset, world_axis)

    def get_plane(self, position):
        normal = self.get_direction(position[:2])
        origin = normal * np.array(
            [self.vs['lx'],
             self.vs['ly'],
             self.vs['lz']], dtype=float) * 0.5

        if len(position) >= 3:
            assert position[2] == 'o'
            origin = np.array([0, 0, 0], dtype=float)

        points = transform_points(
            [origin, origin + normal],
            self.transform
        )
        origin = points[0]
        normal = points[1] - points[0]
        return origin, normal

    def get_axis(self, position, point=None):
        direction = self.get_direction(position)
        if point is None:
            origin = np.array([0, 0, 0], dtype=float)
        else:
            origin = np.asarray(point, dtype=float)

        points = transform_points(
            [origin, origin + direction],
            self.transform
        )

        origin = points[0]
        direction = points[1] - points[0]
        return origin, direction

    def get_mesh(self):
        if self.is_hidden:
            return None

        self.mesh = trimesh.creation.box(
            [self.vs['lx'],
             self.vs['ly'],
             self.vs['lz']],
            transform=self.transform)
        if self.face_colors is not None:
            self.mesh.visual.face_colors = self.face_colors

        return self.mesh